package com.mojo.redisson.config;

import org.springframework.data.redis.core.BoundValueOperations;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.ValueOperations;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.concurrent.TimeUnit;

@Component
public class TokenBucketRateLimiter {

    private static final String REDIS_KEY = "token_bucket_key";
    private static final int MAX_TOKENS = 10; // 桶的容量
    private static final int TOKEN_RATE = 1; // 令牌生成速率，单位 tokens/second
    private static final long TIMEOUT = 10L;
    @Resource
    private RedisTemplate<String, Object> redisTemplate;


    public boolean allowRequest() {
        long currentTime = System.currentTimeMillis() / 1000; // 当前时间戳，秒

        // 获取当前令牌数量
        ValueOperations<String, Object> operations = redisTemplate.opsForValue();
        String tokensStr = (String) operations.get(REDIS_KEY);
        int currentTokens = (tokensStr != null) ? Integer.parseInt(tokensStr) : 0;

        // 计算应该生成的令牌数量
        BoundValueOperations<String, Object> boundValueOps = redisTemplate.boundValueOps(REDIS_KEY);
        Long expire = boundValueOps.getExpire();
        long elapsedSeconds = Math.max(TIMEOUT - (expire == null ? 0 : expire), 0);
        int newTokens = (int) Math.min(MAX_TOKENS, currentTokens + elapsedSeconds * TOKEN_RATE);

        // 更新令牌数量
        operations.setIfAbsent(REDIS_KEY, String.valueOf(newTokens), TIMEOUT, TimeUnit.SECONDS);

        // 判断是否允许新的请求
        return newTokens > 0;
    }
}
